from tqdm import tqdm
import pandas as pd
import numpy as np
from pandas import DataFrame


class DataLoader:
    def __init__(self,chunk_size: int = 100000):
        """
        初始化数据加载器

        :param chunk_size: 读取数据时的分块大小
        """

        self.chunk_size = chunk_size
        self.df: DataFrame | None = None

    def load_data(self, dataset: str or np.ndarray or DataFrame):
        self.df = self._load_data_with_progress(dataset)
        return self.df


    def _load_data_with_progress(self, dataset):
        # 对于NumPy数组
        if isinstance(dataset, np.ndarray):
            return pd.DataFrame(dataset)

        # 对于CSV文件 - 使用chunk_size分块读取
        elif isinstance(dataset, str) and dataset.endswith('.csv'):
            # 先获取总行数用于进度条
            with open(dataset, 'r') as f:
                total = sum(1 for _ in f) - 1  # 减1是去掉header行

            # 分块读取
            chunks = pd.read_csv(dataset, low_memory=False, chunksize=self.chunk_size)
            dfs = []
            for chunk in tqdm(chunks, total=total // self.chunk_size + 1, desc="加载CSV"):
                dfs.append(chunk)
            return pd.concat(dfs, ignore_index=True)

        # 对于DataFrame直接返回
        elif isinstance(dataset, DataFrame):
            return dataset

        # 对于Parquet文件 - 使用pyarrow引擎
        elif isinstance(dataset, str) and dataset.endswith('.parquet'):
            # Parquet文件可以读取元数据获取行数
            import pyarrow.parquet as pq
            pf = pq.ParquetFile(dataset)
            total_rows = pf.metadata.num_rows

            # 分块读取
            batches = pf.iter_batches(batch_size=100000)
            dfs = []
            for batch in tqdm(batches, total=total_rows // 100000 + 1, desc="加载Parquet"):
                dfs.append(batch.to_pandas())
            return pd.concat(dfs, ignore_index=True)

        else:
            raise ValueError("数据类型错误")
